# -*- coding: utf-8 -*-
# pip install llama-index-core

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels

llm = DashScope(
    model_name=DashScopeGenerationModels.QWEN_MAX,
    api_key=os.getenv("AI_DASHSCOPE_API_KEY"),
)

Settings.llm = llm
Settings.embed_model = DashScopeEmbedding(
    model_name="text-embedding-v2",
    api_key=os.getenv("AI_DASHSCOPE_API_KEY"),
)

def test_complete():
    response = llm.complete("帮我推荐一下江浙沪5天的旅游攻略。")
    print(response)
    pass


def test_example():
    # data文件夹下放入文件
    documents = SimpleDirectoryReader("data").load_data()
    index = VectorStoreIndex.from_documents(documents)
    query_engine = index.as_query_engine(llm=llm)
    response = query_engine.query("简单介绍一下新建策略组合的方式？")
    print(response)
    pass


if __name__ == '__main__':
    test_example()

    pass
